"""

"""
import sys

from PyQt5 import QtGui, QtWidgets, QtCore
from PyQt5.QtWidgets import (QApplication, QWidget, QMainWindow, QSlider, QFileDialog, QMessageBox, QHBoxLayout,
                             QLineEdit, QLabel, QVBoxLayout, QPushButton)
from PyQt5.QtGui import (QPainter, QPen, QCursor, QPalette, QColor, QPixmap, QIntValidator)
from PyQt5.QtCore import Qt, pyqtSignal
import operator
from ui.ImageLabel import MyImageLabel
from PIL import Image

from screeninfo import get_monitors


class MyCanvasWindow(QMainWindow):
    closed = pyqtSignal(bool)
    rgbList = pyqtSignal(list)
    closed_ = pyqtSignal()

    def __init__(self, img_path=None, filename=None):
        super().__init__()

        # 彩色列表初始化
        self.colorC = 0
        self.penColor = [QColor(0, 0, 0)]
        self.RGBList = []
        # setMouseTracking设置为False，否则不按下鼠标时也会跟踪鼠标事件
        self.setMouseTracking(False)
        self.setupUi(self)
        self.init_actions()
        self.setStyleSheet("QMainWindow{background-color: white}")

        self.save_path = None
        self.unsaved = False
        self.edited = False
        self.save_path = None

        self.actionUndo.setEnabled(False)
        self.actionRedo.setEnabled(False)

        # set image if invoked from image_window
        self.img_path = img_path
        if img_path and filename:
            self.setWindowTitle(QtCore.QCoreApplication.translate("MainWindow", "Canvas: " + filename))
            self.save_path = img_path
            with Image.open(img_path) as img:
                h = img.height
                w = img.width
            # self.setFixedSize(w, h)
            # self.imageLabel.setFixedSize(w, h)
            # self.imageLabel.setGeometry(QtCore.QRect(0, 0, w, h))
            self.resize_ui(h, w)

    def resize_ui(self, h, w):
        """
        根据图片大小调整窗口及图像标签大小
        Resize the label and window to fit the image size, and show image.
        """
        # get monitor size and resize accordingly
        monitor = get_monitors()[0]
        screen_x = monitor.width
        screen_y = monitor.height
        ratio = w / h
        if h > screen_y:
            x = screen_y * 0.8 * ratio
            y = screen_y * 0.8
        elif w > screen_x:
            x = screen_x * 0.8
            y = screen_x * 0.8 / ratio
        else:
            x = w
            y = h
        self.imageLabel.setGeometry(QtCore.QRect(0, 0, x, y))
        self.imageLabel.setPixmap(QPixmap(self.img_path).scaled(self.imageLabel.size()))
        if sys.platform == 'win32':
            self.setFixedSize(x, y + 25)
        else:
            self.setFixedSize(x, y)

    def resizeEvent(self, event):
        QtWidgets.QMainWindow.resizeEvent(self, event)
        if not self.img_path:
            self.imageLabel.setFixedSize(self.size())

    def closeEvent(self, event):
        # prompt for saving changes
        if self.unsaved:
            msg = QMessageBox()
            msg.setWindowTitle('Warning')
            msg.setText('Do you want to save the changes?')
            msg.setIcon(QMessageBox.Warning)
            msg.setStandardButtons(QMessageBox.Save | QMessageBox.Ignore | QMessageBox.Cancel)
            ret = msg.exec_()
            if ret == QMessageBox.Save:
                self.closed.emit(True)
                self.save()
            elif ret == QMessageBox.Cancel:
                event.ignore()
                return
            else:
                self.closed.emit(False)

        # reset cursor
        while QApplication.overrideCursor() is not None:
            QApplication.restoreOverrideCursor()

        # return to main_window
        self.closed_.emit()

    def setupUi(self, MainWindow):
        MainWindow.setObjectName("Canvas")
        MainWindow.resize(800, 600)
        sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Fixed, QtWidgets.QSizePolicy.Fixed)
        sizePolicy.setHorizontalStretch(0)
        sizePolicy.setVerticalStretch(0)
        sizePolicy.setHeightForWidth(MainWindow.sizePolicy().hasHeightForWidth())
        MainWindow.setSizePolicy(sizePolicy)
        self.centralwidget = QtWidgets.QWidget(MainWindow)
        self.centralwidget.setObjectName("centralwidget")
        MainWindow.setCentralWidget(self.centralwidget)

        # 加bar
        self.menubar = QtWidgets.QMenuBar(MainWindow)
        self.menubar.setGeometry(QtCore.QRect(0, 0, 800, 24))
        self.menubar.setObjectName("menubar")
        self.menuFile = QtWidgets.QMenu(self.menubar)
        self.menuFile.setObjectName("menuFile")
        self.menuEdit = QtWidgets.QMenu(self.menubar)
        self.menuEdit.setObjectName("menuEdit")

        # 加选项
        MainWindow.setMenuBar(self.menubar)
        self.actionSave = QtWidgets.QAction(MainWindow)
        self.actionSave.setObjectName("actionSave")
        self.actionSave_As = QtWidgets.QAction(MainWindow)
        self.actionSave_As.setObjectName("actionSave_As")
        self.actionUndo = QtWidgets.QAction(MainWindow)
        self.actionUndo.setObjectName("actionUndo")
        self.actionRedo = QtWidgets.QAction(MainWindow)
        self.actionRedo.setObjectName("actionRedo")

        # 不知道对不对的
        self.menuMode = QtWidgets.QMenu(self.menubar)
        self.menuMode.setObjectName("mod")
        self.menuStyle = QtWidgets.QMenu(self.menubar)
        self.menuStyle.setObjectName("style")
        self.menuColor = QtWidgets.QMenu(self.menubar)
        self.menuColor.setObjectName("color")
        self.menuSize = QtWidgets.QMenu(self.menubar)
        self.menuSize.setObjectName("size")
        self.actionPen = QtWidgets.QAction(MainWindow)
        self.actionPen.setObjectName("pen")
        self.actionErase = QtWidgets.QAction(MainWindow)
        self.actionErase.setObjectName("erase")
        self.actionErasePaint = QtWidgets.QAction(MainWindow)
        self.actionErasePaint.setObjectName("erase_paint")
        self.actionSetRed = QtWidgets.QAction(MainWindow)
        self.actionSetRed.setObjectName("setRed")
        self.actionSizeAdd = QtWidgets.QAction(MainWindow)
        self.actionSizeAdd.setObjectName("sizeAdd")
        self.actionSizeDec = QtWidgets.QAction(MainWindow)
        self.actionSizeDec.setObjectName("sizeDec")
        self.actionSetGreen = QtWidgets.QAction(MainWindow)
        self.actionSetGreen.setObjectName("setGreen")
        self.actionSetBlack = QtWidgets.QAction(MainWindow)
        self.actionSetBlack.setObjectName("setBlack")
        self.actionSetAddColor = QtWidgets.QAction(MainWindow)
        self.actionSetAddColor.setObjectName("addColor")

        # ImageLabel
        self.imageLabel = MyImageLabel(self.centralwidget, self.penColor)
        self.imageLabel.setGeometry(QtCore.QRect(0, 0, 691, 221))
        self.imageLabel.setObjectName("imageLabel")

        # 逐级别添加
        self.menubar.addAction(self.menuFile.menuAction())
        self.menubar.addAction(self.menuEdit.menuAction())
        self.menubar.addAction(self.menuMode.menuAction())
        self.menubar.addAction(self.menuStyle.menuAction())
        self.menuFile.addAction(self.actionSave)
        self.menuFile.addAction(self.actionSave_As)
        self.menuEdit.addAction(self.actionUndo)
        self.menuEdit.addAction(self.actionRedo)

        self.colorList = [self.actionSetRed, self.actionSetBlack, self.actionSetGreen]

        self.menuStyle.addMenu(self.menuColor)
        self.menuStyle.addMenu(self.menuSize)
        self.menuSize.addActions([self.actionSizeAdd, self.actionSizeDec])
        self.menuMode.addActions([self.actionPen, self.actionErase, self.actionErasePaint])
        self.menuColor.addActions(self.colorList)
        self.menuColor.addAction(self.actionSetAddColor)

        self.retranslateUi(MainWindow)
        QtCore.QMetaObject.connectSlotsByName(MainWindow)

    def retranslateUi(self, MainWindow):
        _translate = QtCore.QCoreApplication.translate
        MainWindow.setWindowTitle(_translate("MainWindow", "Canvas"))
        self.menuFile.setTitle(_translate("MainWindow", "File"))
        self.menuEdit.setTitle(_translate("MainWindow", "Edit"))
        self.actionSave.setText(_translate("MainWindow", "Save"))
        self.actionSave.setShortcut(_translate("MainWindow", "Ctrl+S"))
        self.actionSave_As.setText(_translate("MainWindow", "Save As..."))
        self.actionSave_As.setShortcut(_translate("MainWindow", "Ctrl+Shift+S"))
        self.actionUndo.setText(_translate("MainWindow", "Undo"))
        self.actionUndo.setShortcut(_translate("MainWindow", "Ctrl+Z"))
        self.actionRedo.setText(_translate("MainWindow", "Redo"))
        self.actionRedo.setShortcut(_translate("MainWindow", "Ctrl+Shift+Z"))

        self.menuStyle.setTitle(_translate("MainWindow", "Style"))
        self.menuMode.setTitle(_translate("MainWindow", "Mode"))
        self.menuColor.setTitle(_translate("MainWindow", "Color"))
        self.menuSize.setTitle(_translate("MainWindow", "Size"))
        self.actionPen.setText(_translate("MainWindow", "Pen"))
        self.actionPen.setShortcut(_translate("MainWindow", "Ctrl+P"))
        self.actionErase.setText(_translate("MainWindow", "Erase"))
        self.actionErase.setShortcut(_translate("MainWindow", "Ctrl+E"))
        self.actionErasePaint.setText(_translate("MainWindow", "Erase Paint"))
        self.actionErasePaint.setShortcut(_translate("MainWindow", "Ctrl+Shift+E"))
        self.actionSizeAdd.setText(_translate("MainWindow", "Size+"))
        self.actionSizeAdd.setShortcut(_translate("MainWindow", "Ctrl+="))
        self.actionSizeDec.setText(_translate("MainWindow", "Size-"))
        self.actionSizeDec.setShortcut(_translate("MainWindow", "Ctrl+-"))
        self.actionSetRed.setText(_translate("MainWindow", "Red"))
        self.actionSetRed.setShortcut(_translate("MainWindow", "Ctrl+Alt+R"))
        self.actionSetGreen.setText(_translate("MainWindow", "Green"))
        self.actionSetGreen.setShortcut(_translate("MainWindow", "Ctrl+Alt+G"))
        self.actionSetBlack.setText(_translate("MainWindow", "Black"))
        self.actionSetBlack.setShortcut(_translate("MainWindow", "Ctrl+Alt+B"))
        self.actionSetAddColor.setText(_translate("MainWindow", "Add Color"))
        self.actionSetAddColor.setShortcut(_translate("MainWindow", "Ctrl+Alt+="))

    def init_actions(self):
        self.actionSetBlack.triggered.connect(self.setColorBlack)
        self.actionSetRed.triggered.connect(self.setColorRed)
        self.actionSetGreen.triggered.connect(self.setColorGreen)
        self.actionSave.triggered.connect(self.get_screenshot)
        self.actionSizeAdd.triggered.connect(self.imageLabel.sizeAdd)
        self.actionSizeDec.triggered.connect(self.imageLabel.sizeDec)
        self.actionPen.triggered.connect(self.set_pen)
        self.actionErase.triggered.connect(self.set_eraser)
        self.actionErasePaint.triggered.connect(self.set_eraser_paint)
        self.actionUndo.triggered.connect(self.imageLabel.undo)
        self.actionRedo.triggered.connect(self.imageLabel.redo)
        self.actionSave.triggered.connect(self.save)
        self.actionSave_As.triggered.connect(self.save_as)
        self.actionSetAddColor.triggered.connect(self.open_add_color_window)
        self.imageLabel.undo_signal.connect(lambda x: self.actionUndo.setEnabled(x))
        self.imageLabel.redo_signal.connect(lambda x: self.actionRedo.setEnabled(x))
        self.imageLabel.unsaved_signal.connect(lambda: self.set_unsaved(True))
        self.imageLabel.edited_signal.connect(lambda x: self.set_edited(x))

    def set_edited(self, x):
        self.edited = x

    def set_unsaved(self, x):
        self.unsaved = x

    def set_pen(self):
        self.menuColor.setEnabled(True)
        self.imageLabel.setPen()

    def set_eraser(self):
        self.menuColor.setEnabled(False)
        self.imageLabel.setErase()

    def set_eraser_paint(self):
        self.menuColor.setEnabled(False)
        self.imageLabel.setErasePaint()

    def save(self):
        if not self.save_path:
            self.save_as()
        else:
            self.get_screenshot().save(self.save_path, 'PNG')
            self.set_unsaved(False)

    def save_as(self):
        self.save_path, _ = QFileDialog.getSaveFileName(self, 'Save As PNG Image', self.save_path, 'PNG Image (*.png)')
        if self.save_path:
            img = self.get_screenshot()
            img.save(self.save_path, 'PNG')
        else:
            msg = QMessageBox()
            msg.setWindowTitle('Error')
            msg.setText('File could not be saved.')
            msg.setIcon(QMessageBox.Critical)
            _ = msg.exec_()

    def setColor(self, color):
        self.penColor[0] = QColor(color[0], color[1], color[2])

    def setColorRed(self):
        self.setColor([255, 0, 0])

    def setColorBlack(self):
        self.setColor([0, 0, 0])

    def setColorGreen(self):
        self.setColor([0, 255, 0])

    def get_screenshot(self):
        screen = QApplication.primaryScreen()
        return screen.grabWindow(self.centralWidget().winId()).toImage()

    def open_add_color_window(self):
        self.win = AddColorWindow(self.RGBList)
        self.win.setWindowModality(QtCore.Qt.ApplicationModal)  # freeze window when widget is alive
        self.win.show()
        # self.actionSetAddColor.triggered.connect(lambda: self.imageLabel.setColor(self.RBGList))
        self.win.bottom.clicked.connect(self.add_color)

    def add_color(self):
        _translate = QtCore.QCoreApplication.translate
        self.colorC += 1
        count = self.colorC
        self.RGBList.append([int(self.win.R.text()), int(self.win.G.text()), int(self.win.B.text())])
        actionSetTemp = QtWidgets.QAction(self)
        actionSetTemp.setObjectName("setColor" + str(len(self.colorList) + 1))
        actionSetTemp.setText(_translate("MainWindow", "Color" + str(self.colorC)+str(self.RGBList[count-1])))
        actionSetTemp.setShortcut(_translate("MainWindow", "Ctrl+Alt+"+str(count)))
        self.colorList.append(actionSetTemp)
        self.menuColor.addActions(self.colorList)
        self.menuColor.addAction(self.actionSetAddColor)
        self.win.R.text()
        actionSetTemp.triggered.connect(lambda: self.setColor(self.RGBList[count-1]))
        self.setColor(self.RGBList[count - 1])
        self.win.close()


class AddColorWindow(QWidget):
    add_color_signal = pyqtSignal()

    def __init__(self, RGBList):
        super(AddColorWindow, self).__init__()
        self.RGBList = RGBList
        self.setFixedSize(300, 200)
        self.setWindowTitle("Add Color")
        QtCore.QMetaObject.connectSlotsByName(self)
        self.setUp_ui()

    def setUp_ui(self):
        hbox = QHBoxLayout()
        validator = QIntValidator()
        validator.setRange(0, 255)
        self.R = QLineEdit()
        self.R.setValidator(validator)
        self.R.setText('0')
        self.G = QLineEdit()
        self.G.setValidator(validator)
        self.G.setText('0')
        self.B = QLineEdit()
        self.B.setValidator(validator)
        self.B.setText('0')
        self.RText = QLabel(self)  # 静态标签
        self.RText.setText(u"R:")
        self.GText = QLabel(self)  # 静态标签
        self.GText.setText(u"G:")
        self.BText = QLabel(self)  # 静态标签
        self.BText.setText(u"B:")
        self.bottom = QPushButton("save", self)

        hbox.addWidget(self.RText, 1)
        hbox.addWidget(self.R, 1)
        hbox.addWidget(self.GText, 1)
        hbox.addWidget(self.G, 1)
        hbox.addWidget(self.BText, 1)
        hbox.addWidget(self.B, 1)

        vbox = QVBoxLayout()
        vbox.addLayout(hbox, 1)
        vbox.addWidget(self.bottom, 1)
        self.setLayout(vbox)


if __name__ == '__main__':
    """
    程序入口，新建一个主窗口并打开
    """
    app = QApplication(sys.argv)
    win = MyCanvasWindow('./sample_img/cat.bmp', 'cat.bmp')
    # win = AddColorWindow([])
    win.show()
    sys.exit(app.exec_())
